#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PCAP File Parser Script
PCAP文件解析脚本

功能：
1. 解析pcap文件，提取网络包信息
2. 识别TCP、UDP、ICMP、DNS等协议
3. 生成结构化的JSON数据
4. 支持批量处理

作者：AI Assistant
版本：1.0
日期：2025-08-15
"""

import sys
import json
import argparse
from pathlib import Path
from datetime import datetime
from collections import defaultdict

try:
    from scapy.all import rdpcap, IP, TCP, UDP, ICMP, DNS
    SCAPY_AVAILABLE = True
except ImportError:
    SCAPY_AVAILABLE = False
    print("Warning: scapy library not installed, will use basic parsing mode")

class PCAPParser:
    def __init__(self):
        self.protocols = {
            1: "ICMP",
            6: "TCP", 
            17: "UDP",
            53: "DNS"
        }
        
    def parse_pcap_file(self, pcap_file_path, output_dir):
        """解析pcap文件并生成JSON"""
        pcap_file = Path(pcap_file_path)
        output_dir = Path(output_dir)
        
        if not pcap_file.exists():
            print(f"Pcap file does not exist: {pcap_file}")
            return False
            
        if not output_dir.exists():
            output_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"Starting pcap file parsing: {pcap_file.name}")
        
        try:
            if SCAPY_AVAILABLE:
                packets_data = self._parse_with_scapy(pcap_file)
            else:
                packets_data = self._parse_basic(pcap_file)
            
            # 生成输出文件名
            output_file = output_dir / f"{pcap_file.stem}.json"
            
            # 保存JSON文件
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(packets_data, f, ensure_ascii=False, indent=2)
            
            print(f"JSON file generated successfully: {output_file}")
            print(f"Parsed packets: {len(packets_data['packets'])}")
            
            return True
            
        except Exception as e:
            print(f"Error parsing pcap file: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def _parse_with_scapy(self, pcap_file):
        """使用scapy解析pcap文件"""
        print("Using scapy for advanced parsing...")
        
        # 读取pcap文件
        packets = rdpcap(str(pcap_file))
        
        # 初始化数据结构
        packets_data = {
            "metadata": {
                "filename": pcap_file.name,
                "parser_version": "1.0",
                "parse_timestamp": datetime.now().isoformat(),
                "total_packets": len(packets),
                "statistics": {
                    "tcp_packets": 0,
                    "udp_packets": 0,
                    "icmp_packets": 0,
                    "dns_packets": 0,
                    "other_packets": 0
                }
            },
            "packets": []
        }
        
        # 解析每个包
        for i, packet in enumerate(packets):
            packet_info = self._extract_packet_info(packet, i + 1)
            packets_data["packets"].append(packet_info)
            
            # 更新统计信息
            self._update_statistics(packets_data["metadata"]["statistics"], packet_info)
        
        return packets_data
    
    def _parse_basic(self, pcap_file):
        """基础解析模式（当scapy不可用时）"""
        print("Using basic parsing mode...")
        
        packets_data = {
            "metadata": {
                "filename": pcap_file.name,
                "parser_version": "1.0",
                "parse_timestamp": datetime.now().isoformat(),
                "total_packets": 0,
                "note": "Basic parsing mode, limited information",
                "statistics": {
                    "tcp_packets": 0,
                    "udp_packets": 0,
                    "icmp_packets": 0,
                    "dns_packets": 0,
                    "other_packets": 0
                }
            },
            "packets": []
        }
        
        print("Basic parsing mode has limited functionality, recommend installing scapy library")
        return packets_data
    
    def _extract_packet_info(self, packet, packet_number):
        """提取单个包的信息"""
        packet_info = {
            "packet_number": packet_number,
            "timestamp": float(packet.time) if hasattr(packet, 'time') else 0,
            "datetime": datetime.fromtimestamp(float(packet.time)).isoformat() if hasattr(packet, 'time') else "",
            "length": len(packet),
            "protocol": "UNKNOWN"
        }
        
        # 提取IP层信息
        if IP in packet:
            ip_layer = packet[IP]
            packet_info.update({
                "src_ip": ip_layer.src,
                "dst_ip": ip_layer.dst,
                "protocol_num": ip_layer.proto,
                "ttl": ip_layer.ttl
            })
            
            # 识别传输层协议
            if TCP in packet:
                tcp_layer = packet[TCP]
                packet_info.update({
                    "src_port": tcp_layer.sport,
                    "dst_port": tcp_layer.dport,
                    "transport": "TCP",
                    "protocol": "TCP",
                    "tcp_flags": self._extract_tcp_flags(tcp_layer),
                    "sequence": tcp_layer.seq,
                    "acknowledgment": tcp_layer.ack,
                    "window_size": tcp_layer.window
                })
                
            elif UDP in packet:
                udp_layer = packet[UDP]
                packet_info.update({
                    "src_port": udp_layer.sport,
                    "dst_port": udp_layer.dport,
                    "transport": "UDP",
                    "protocol": "UDP"
                })
                
                # 检查是否为DNS
                if DNS in packet:
                    packet_info["protocol"] = "DNS"
                    packet_info["dns_info"] = self._extract_dns_info(packet[DNS])
                    
            elif ICMP in packet:
                icmp_layer = packet[ICMP]
                packet_info.update({
                    "transport": "ICMP",
                    "protocol": "ICMP",
                    "icmp_type": icmp_layer.type,
                    "icmp_code": icmp_layer.code
                })
                
            else:
                # 其他IP协议
                packet_info["protocol"] = self.protocols.get(ip_layer.proto, f"PROTO_{ip_layer.proto}")
                
        else:
            # 非IP包
            packet_info.update({
                "src_ip": "N/A",
                "dst_ip": "N/A",
                "protocol_num": 0,
                "src_port": 0,
                "dst_port": 0,
                "transport": "UNKNOWN"
            })
        
        return packet_info
    
    def _extract_tcp_flags(self, tcp_layer):
        """提取TCP标志位"""
        flags = []
        if tcp_layer.flags & 0x01:  # FIN
            flags.append("FIN")
        if tcp_layer.flags & 0x02:  # SYN
            flags.append("SYN")
        if tcp_layer.flags & 0x04:  # RST
            flags.append("RST")
        if tcp_layer.flags & 0x08:  # PSH
            flags.append("PSH")
        if tcp_layer.flags & 0x10:  # ACK
            flags.append("ACK")
        if tcp_layer.flags & 0x20:  # URG
            flags.append("URG")
        return flags
    
    def _extract_dns_info(self, dns_layer):
        """提取DNS信息"""
        dns_info = {
            "qr": "Query" if dns_layer.qr == 0 else "Response",
            "opcode": dns_layer.opcode,
            "questions": dns_layer.qdcount,
            "answers": dns_layer.ancount
        }
        
        # 提取查询信息
        if hasattr(dns_layer, 'qd') and dns_layer.qd:
            dns_info["query"] = str(dns_layer.qd.qname)
            dns_info["qtype"] = dns_layer.qd.qtype
            
        return dns_info
    
    def _update_statistics(self, stats, packet_info):
        """更新统计信息"""
        protocol = packet_info.get("protocol", "UNKNOWN")
        
        if protocol == "TCP":
            stats["tcp_packets"] += 1
        elif protocol == "UDP":
            stats["udp_packets"] += 1
        elif protocol == "ICMP":
            stats["icmp_packets"] += 1
        elif protocol == "DNS":
            stats["dns_packets"] += 1
        else:
            stats["other_packets"] += 1

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="PCAP File Parser Tool")
    parser.add_argument("pcap_file", help="Input pcap file path")
    parser.add_argument("output_dir", help="Output JSON file directory")
    
    args = parser.parse_args()
    
    # 创建解析器实例
    pcap_parser = PCAPParser()
    
    # 解析pcap文件
    success = pcap_parser.parse_pcap_file(args.pcap_file, args.output_dir)
    
    if success:
        print("PCAP parsing completed!")
        sys.exit(0)
    else:
        print("PCAP parsing failed!")
        sys.exit(1)

if __name__ == "__main__":
    main() 